/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/* !
 * \file softmax_flashv2_base_impl.h
 * \brief
 */
#ifndef IMPL_ACTIVATION_SOFTMAX_SOFTMAX_FLASHV2_BASE_IMPL_H
#define IMPL_ACTIVATION_SOFTMAX_SOFTMAX_FLASHV2_BASE_IMPL_H

#if __CCE_AICORE__ == 300
#include "v300/softmax_flashv2_impl.h"
#elif __CCE_AICORE__ == 220
#include "v220/softmax_flashv2_impl.h"
#elif __CCE_AICORE__ == 200
#include "v200/softmax_flashv2_impl.h"
#endif

namespace AscendC {
__aicore__ inline constexpr SoftMaxTiling SoftMaxFlashV2TilingFuncImpl(const uint32_t srcM, const uint32_t srcK,
    const uint32_t dataTypeSize1, const uint32_t dataTypeSize2, const uint32_t localWorkSpaceSize,
    const bool isUpdate = false, const bool isBasicBlock = false, const bool isDataFormatNZ = false)
{
    SoftMaxTiling softmaxTiling;
    const uint32_t elementNumPerBlk = ONE_BLK_SIZE / dataTypeSize2;
    softmaxTiling.srcM = srcM;
    softmaxTiling.srcK = srcK;
    softmaxTiling.srcSize = srcM * srcK;

    softmaxTiling.outMaxM = srcM;
    softmaxTiling.outMaxK = elementNumPerBlk;
    softmaxTiling.outMaxSize = srcM * elementNumPerBlk;

    if (isDataFormatNZ) {
        softmaxTiling.reduceM = localWorkSpaceSize / (SOFTMAX_SHAPE_NZ_BASIC_COUNT * HALF_FACTOR + srcK);
    } else {
        if (isBasicBlock && srcK % FLOAT_REPEAT_SIZE == 0 && srcM % SOFTMAX_BASIC_TILE_NUM == 0) {
            softmaxTiling.reduceM =
                CalculateNDSplitM(localWorkSpaceSize, dataTypeSize1, elementNumPerBlk, { srcM, srcK }, isBasicBlock);
        } else {
            softmaxTiling.reduceM = (dataTypeSize1 == B16_BYTE_SIZE) ?
                localWorkSpaceSize / (SOFTMAX_COMPUTE_DIM * elementNumPerBlk + srcK + FLOAT_REPEAT_SIZE) :
                localWorkSpaceSize / (elementNumPerBlk + FLOAT_REPEAT_SIZE);
        }
    }

    if (softmaxTiling.reduceM < srcM && softmaxTiling.reduceM > SOFTMAX_BASIC_TILE_NUM) {
        softmaxTiling.reduceM = softmaxTiling.reduceM / SOFTMAX_BASIC_TILE_NUM * SOFTMAX_BASIC_TILE_NUM;
    }
    softmaxTiling.reduceM = softmaxTiling.reduceM < srcM ? softmaxTiling.reduceM : srcM;

    softmaxTiling.reduceK = elementNumPerBlk;
    softmaxTiling.reduceSize = softmaxTiling.reduceM * elementNumPerBlk;

    softmaxTiling.splitM = softmaxTiling.reduceM;
    softmaxTiling.splitK = srcK;
    softmaxTiling.splitSize = softmaxTiling.reduceM * srcK;

    softmaxTiling.rangeM = srcM / softmaxTiling.reduceM;
    softmaxTiling.tailM = srcM % softmaxTiling.reduceM;

    softmaxTiling.tailSplitSize = softmaxTiling.tailM * srcK;
    softmaxTiling.tailReduceSize = softmaxTiling.tailM * elementNumPerBlk;
    return softmaxTiling;
}

template <typename T1, typename T2, bool isUpdate, bool isBasicBlock, bool isDataFormatNZ, const SoftmaxConfig& config>
__aicore__ inline void SoftmaxFlashV2Impl(const LocalTensor<T1>& dstTensor, const LocalTensor<T2>& sumTensor,
    const LocalTensor<T2>& maxTensor, const LocalTensor<T1>& srcTensor, const LocalTensor<T1>& expMaxTensor,
    const LocalTensor<T2>& inSumTensor, const LocalTensor<T2>& inMaxTensor, const LocalTensor<float>& workLocal,
    const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo)
{
    SetMaskNorm();
    ResetMask();
    ShapeInfo srcShape = srcTensor.GetShapeInfo();
    uint32_t workLocalSize = workLocal.GetSize();
    LastAxisShapeND srcNDinfo;
    LastAxisShapeND originalSrcShape;
    if (softmaxShapeInfo.srcM == 0 || softmaxShapeInfo.srcK == 0) {
        srcNDinfo = GetLastAxisShapeND(srcShape);
        originalSrcShape = GetLastAxisOriginShapeND(srcShape);
    } else {
        srcNDinfo = { softmaxShapeInfo.srcM, softmaxShapeInfo.srcK };
        originalSrcShape = { softmaxShapeInfo.oriSrcM, softmaxShapeInfo.oriSrcK };
    }
    if constexpr (isDataFormatNZ) {
        if constexpr (config.isCheckTiling) {
            if (unlikely(srcNDinfo.k != tiling.srcK || srcNDinfo.m != tiling.srcM)) {
                SoftMaxTiling newTiling = SoftMaxFlashV2TilingFuncImpl(srcNDinfo.m, srcNDinfo.k, sizeof(T1), sizeof(T2),
                    workLocalSize, isUpdate, false, true);
                SoftMaxFlashV2NZImpl<T1, T2, isUpdate, isBasicBlock>(dstTensor, sumTensor, maxTensor, srcTensor,
                    expMaxTensor, inSumTensor, inMaxTensor, workLocal, originalSrcShape, newTiling);
            } else {
                SoftMaxFlashV2NZImpl<T1, T2, isUpdate, isBasicBlock>(dstTensor, sumTensor, maxTensor, srcTensor,
                    expMaxTensor, inSumTensor, inMaxTensor, workLocal, originalSrcShape, tiling);
            }
        } else {
            SoftMaxFlashV2NZImpl<T1, T2, isUpdate, isBasicBlock>(dstTensor, sumTensor, maxTensor, srcTensor,
                expMaxTensor, inSumTensor, inMaxTensor, workLocal, originalSrcShape, tiling);
        }
    } else {
        if constexpr (config.isCheckTiling) {
            if (unlikely(srcNDinfo.k != tiling.srcK || srcNDinfo.m != tiling.srcM)) {
                SoftMaxTiling newTiling = SoftMaxFlashV2TilingFuncImpl(srcNDinfo.m, srcNDinfo.k, sizeof(T1), sizeof(T2),
                    workLocalSize, isUpdate, isBasicBlock);
                SoftmaxFlashV2PostProcess<T1, T2, isUpdate, isBasicBlock>(dstTensor, sumTensor, maxTensor, srcTensor,
                    expMaxTensor, inSumTensor, inMaxTensor, workLocal, originalSrcShape, newTiling);
            } else {
                SoftmaxFlashV2PostProcess<T1, T2, isUpdate, isBasicBlock>(dstTensor, sumTensor, maxTensor, srcTensor,
                    expMaxTensor, inSumTensor, inMaxTensor, workLocal, originalSrcShape, tiling);
            }
        } else {
            SoftmaxFlashV2PostProcess<T1, T2, isUpdate, isBasicBlock>(dstTensor, sumTensor, maxTensor, srcTensor,
                expMaxTensor, inSumTensor, inMaxTensor, workLocal, originalSrcShape, tiling);
        }
    }
}

template <typename T1, typename T2, bool isUpdate, bool isBasicBlock, bool isDataFormatNZ, const SoftmaxConfig& config>
__aicore__ inline void SoftmaxFlashV2Impl(const LocalTensor<T1>& dstTensor, const LocalTensor<T2>& sumTensor,
    const LocalTensor<T2>& maxTensor, const LocalTensor<T1>& srcTensor, const LocalTensor<T1>& expMaxTensor,
    const LocalTensor<T2>& inSumTensor, const LocalTensor<T2>& inMaxTensor, const SoftMaxTiling& tiling,
    const SoftMaxShapeInfo& softmaxShapeInfo)
{
    LocalTensor<float> workLocal;
    PopStackBuffer<float, TPosition::LCM>(workLocal);
    SoftmaxFlashV2Impl<T1, T2, isUpdate, isBasicBlock, isDataFormatNZ, config>(dstTensor, sumTensor, maxTensor,
        srcTensor, expMaxTensor, inSumTensor, inMaxTensor, workLocal, tiling, softmaxShapeInfo);
}

template <typename T1, typename T2, bool isUpdate, bool isBasicBlock, bool isDataFormatNZ, const SoftmaxConfig& config>
__aicore__ inline void SoftmaxFlashV2Impl(const LocalTensor<T1>& dstTensor, const LocalTensor<T2>& sumTensor,
    const LocalTensor<T2>& maxTensor, const LocalTensor<T1>& srcTensor, const LocalTensor<T1>& expMaxTensor,
    const LocalTensor<T2>& inSumTensor, const LocalTensor<T2>& inMaxTensor,
    const LocalTensor<uint8_t>& sharedTmpBuffer, const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo)
{
    auto workLocal = sharedTmpBuffer.ReinterpretCast<float>();
    SoftmaxFlashV2Impl<T1, T2, isUpdate, isBasicBlock, isDataFormatNZ, config>(dstTensor, sumTensor, maxTensor,
        srcTensor, expMaxTensor, inSumTensor, inMaxTensor, workLocal, tiling, softmaxShapeInfo);
}
}
#endif // IMPL_ACTIVATION_SOFTMAX_SOFTMAX_FLASHV2_BASE_IMPL_H